﻿/// Copyright © 2013 Microsoft Corporation. 
/// This project is released under the Microsoft Public License (MS-PL)
///
namespace DataMorph.Generator
{
    using System;
    using System.Collections.Generic;
    using System.Linq;

    public static class EnumerableExtensions
    {
         public static IEnumerable<IEnumerable<TSource>> Batch<TSource>(this IEnumerable<TSource> source, int batchCount)
         {
             return Batch(source, batchCount, x => x);
         }

         public static IEnumerable<TResult> Batch<TSource, TResult>(
                                                            this IEnumerable<TSource> source, 
                                                            int batchCount,
                                                            Func<IEnumerable<TSource>, TResult> resultSelector)
         {
             var remainder = 0d;
             var itemCount = 0;
             var sourceItemCount = source.Count();
             var itemsPerBatch = sourceItemCount / (double)batchCount;

             for (var i = 0; i < batchCount; i++)
             {
                 var itemsInBatch = Convert.ToInt32(Math.Floor(itemsPerBatch + remainder));
                 remainder = itemsPerBatch + remainder - itemsInBatch;

                 if (i == (batchCount - 1))
                 {
                     itemsInBatch = sourceItemCount - itemCount;
                 }

                 yield return resultSelector(source.Skip(itemCount).Take(itemsInBatch));

                 itemCount += itemsInBatch;
             }
         }

         public static IEnumerable<TResult> Batch<TResult>(int sourceItemCount, int batchCount, Func<int, int, TResult> resultSelector)
         {
             var remainder = 0d;
             var itemCount = 0;
             var itemsPerBatch = sourceItemCount / (double)batchCount;

             for (var i = 0; i < batchCount; i++)
             {
                 var itemsInBatch = Convert.ToInt32(Math.Floor(itemsPerBatch + remainder));
                 remainder = itemsPerBatch + remainder - itemsInBatch;

                 if (i == (batchCount - 1))
                 {
                     itemsInBatch = sourceItemCount - itemCount;
                 }

                 yield return resultSelector(i, itemsInBatch);

                 itemCount += itemsInBatch;
             }
         }
    }
}